{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Context\n",
    "\n",
    "A model context supports storage and retrieval of Chat Completion messages.\n",
    "It is always used together with a model client to generate LLM-based responses.\n",
    "\n",
    "For example, {py:mod}`~autogen_core.model_context.BufferedChatCompletionContext`\n",
    "is a most-recent-used (MRU) context that stores the most recent `buffer_size`\n",
    "number of messages. This is useful to avoid context overflow in many LLMs.\n",
    "\n",
    "Let's see an example that uses\n",
    "{py:mod}`~autogen_core.model_context.BufferedChatCompletionContext`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "from autogen_core import AgentId, MessageContext, RoutedAgent, SingleThreadedAgentRuntime, message_handler\n",
    "from autogen_core.model_context import BufferedChatCompletionContext\n",
    "from autogen_core.models import AssistantMessage, ChatCompletionClient, SystemMessage, UserMessage\n",
    "from autogen_ext.models.openai import OpenAIChatCompletionClient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Message:\n",
    "    content: str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleAgentWithContext(RoutedAgent):\n",
    "    def __init__(self, model_client: ChatCompletionClient) -> None:\n",
    "        super().__init__(\"A simple agent\")\n",
    "        self._system_messages = [SystemMessage(content=\"You are a helpful AI assistant.\")]\n",
    "        self._model_client = model_client\n",
    "        self._model_context = BufferedChatCompletionContext(buffer_size=5)\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_user_message(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        # Prepare input to the chat completion model.\n",
    "        user_message = UserMessage(content=message.content, source=\"user\")\n",
    "        # Add message to model context.\n",
    "        await self._model_context.add_message(user_message)\n",
    "        # Generate a response.\n",
    "        response = await self._model_client.create(\n",
    "            self._system_messages + (await self._model_context.get_messages()),\n",
    "            cancellation_token=ctx.cancellation_token,\n",
    "        )\n",
    "        # Return with the model's response.\n",
    "        assert isinstance(response.content, str)\n",
    "        # Add message to model context.\n",
    "        await self._model_context.add_message(AssistantMessage(content=response.content, source=self.metadata[\"type\"]))\n",
    "        return Message(content=response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's try to ask follow up questions after the first one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: Hello, what are some fun things to do in Seattle?\n",
      "Response: Seattle offers a variety of fun activities and attractions. Here are some highlights:\n",
      "\n",
      "1. **Pike Place Market**: Visit this iconic market to explore local vendors, fresh produce, artisanal products, and watch the famous fish throwing.\n",
      "\n",
      "2. **Space Needle**: Take a trip to the observation deck for stunning panoramic views of the city, Puget Sound, and the surrounding mountains.\n",
      "\n",
      "3. **Chihuly Garden and Glass**: Marvel at the stunning glass art installations created by artist Dale Chihuly, located right next to the Space Needle.\n",
      "\n",
      "4. **Seattle Waterfront**: Enjoy a stroll along the waterfront, visit the Seattle Aquarium, and take a ferry ride to nearby islands like Bainbridge Island.\n",
      "\n",
      "5. **Museum of Pop Culture (MoPOP)**: Explore exhibits on music, science fiction, and pop culture in this architecturally striking building.\n",
      "\n",
      "6. **Seattle Art Museum (SAM)**: Discover an extensive collection of art from around the world, including contemporary and Native American art.\n",
      "\n",
      "7. **Gas Works Park**: Relax in this unique park that features remnants of an old gasification plant, offering great views of the Seattle skyline and Lake Union.\n",
      "\n",
      "8. **Discovery Park**: Enjoy nature trails, beaches, and beautiful views of the Puget Sound and the Olympic Mountains in this large urban park.\n",
      "\n",
      "9. **Ballard Locks**: Watch boats navigate the locks and see fish swimming upstream during the salmon migration season.\n",
      "\n",
      "10. **Fremont Troll**: Check out this quirky public art installation under a bridge in the Fremont neighborhood.\n",
      "\n",
      "11. **Underground Tour**: Take an entertaining guided tour through the underground passages of Pioneer Square to learn about Seattle's history.\n",
      "\n",
      "12. **Brewery Tours**: Seattle is known for its craft beer scene. Visit local breweries for tastings and tours.\n",
      "\n",
      "13. **Seattle Center**: Explore the cultural complex that includes the Space Needle, MoPOP, and various festivals and events throughout the year.\n",
      "\n",
      "These are just a few options, and Seattle has something for everyone, whether you're into outdoor activities, culture, history, or food!\n",
      "-----\n",
      "Question: What was the first thing you mentioned?\n",
      "Response: The first thing I mentioned was **Pike Place Market**. It's an iconic market in Seattle known for its local vendors, fresh produce, artisanal products, and the famous fish throwing by the fishmongers. It's a vibrant place full of sights, sounds, and delicious food.\n"
     ]
    }
   ],
   "source": [
    "model_client = OpenAIChatCompletionClient(\n",
    "    model=\"gpt-4o-mini\",\n",
    "    # api_key=\"sk-...\", # Optional if you have an OPENAI_API_KEY set in the environment.\n",
    ")\n",
    "\n",
    "runtime = SingleThreadedAgentRuntime()\n",
    "await SimpleAgentWithContext.register(\n",
    "    runtime,\n",
    "    \"simple_agent_context\",\n",
    "    lambda: SimpleAgentWithContext(model_client=model_client),\n",
    ")\n",
    "# Start the runtime processing messages.\n",
    "runtime.start()\n",
    "agent_id = AgentId(\"simple_agent_context\", \"default\")\n",
    "\n",
    "# First question.\n",
    "message = Message(\"Hello, what are some fun things to do in Seattle?\")\n",
    "print(f\"Question: {message.content}\")\n",
    "response = await runtime.send_message(message, agent_id)\n",
    "print(f\"Response: {response.content}\")\n",
    "print(\"-----\")\n",
    "\n",
    "# Second question.\n",
    "message = Message(\"What was the first thing you mentioned?\")\n",
    "print(f\"Question: {message.content}\")\n",
    "response = await runtime.send_message(message, agent_id)\n",
    "print(f\"Response: {response.content}\")\n",
    "\n",
    "# Stop the runtime processing messages.\n",
    "await runtime.stop()\n",
    "await model_client.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the second response, you can see the agent now can recall its own previous responses."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
